import numpy as np
import torch

from dataLoader import _load


def _numpy_to_cuda(x):
    x = torch.from_numpy(x)
    if torch.cuda.is_available():
        if x.is_cuda:
            return x
        else:
            return x.cuda()
    else:
        return x


_to_tensor = _numpy_to_cuda  # gpu
# 相对路径是相对于_load而言的
data_path = 'dataset/'
bfm_path = 'bfm/generate/'
# w_shp = _load(bfm_path + 'w_shp.pkl')
# w_exp = _load(bfm_path + 'w_exp.pkl')
# u = _load(bfm_path + 'u.pkl')
w_shp = _load(bfm_path + 'w_shp_f_s.pkl')
w_exp = _load(bfm_path + 'w_exp_f_s.pkl')
u = _load(bfm_path + 'u_s.pkl')
# key_point = _load(bfm_path + 'key_point.pkl')
# key_point = _load(bfm_path+'keypoints_sim.pkl')
key_point = _load(bfm_path+'keypoints_s.pkl')
# mean_std = _load(data_path + 'mean_std_fixed.pkl')
mean_std = _load(data_path + 'mean_std.pkl')
w = np.concatenate((w_shp, w_exp), axis=1)
w_base = w[key_point]
w_norm = np.linalg.norm(w, axis=0)
dim_shp = 199
dim_exp = 29
